import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;

public class Enigma {

    // 定义轮子
    private static int[] SlowLeft = {24, 25, 26, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
    private static int[] SlowRight = {21, 3, 15, 1, 19, 10, 14, 26, 20, 8, 16, 7, 22, 4, 11, 5, 17, 9, 12, 23, 18, 2, 25, 6, 24, 13};
    private static int[] MidLeft = {26, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25};
    private static int[] MidRight = {20, 1, 6, 4, 15, 3, 14, 12, 23, 5, 16, 2, 22, 19, 11, 18, 25, 24, 13, 7, 10, 8, 21, 9, 26, 17};
    private static int[] FastLeft = {26, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25};
    private static int[] FastRight = {14, 8, 18, 26, 17, 20, 22, 10, 3, 13, 11, 4, 23, 5, 24, 9, 12, 25, 16, 19, 6, 15, 21, 2, 7, 1};

    private static final char[] Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ".toCharArray();

    // 轮子的转动
    public static void rotateWheel(int[] wheels) {
        int temp = wheels[25];
        System.arraycopy(wheels, 0, wheels, 1, 25);
        wheels[0] = temp;
    }

    // 查找元素
    public static int searchElement(char[] wheels, char element) {
        for (int i = 0; i < wheels.length; i++) {
            if (wheels[i] == element) {
                return i;
            }
        }
        System.out.println("查询失败");
        return -1;
    }

    // 查找数字
    public static int searchNumber(int[] wheels, int number) {
        for (int i = 0; i < wheels.length; i++) {
            if (wheels[i] == number) {
                return i;
            }
        }
        System.out.println("查询失败");
        return -1;
    }

    // 加密函数，返回加密后的字符串
    public static String encryptFunction(String message) {
        StringBuilder encryptedMessage = new StringBuilder();
        int count1 = 0;
        int count2 = 0;

        for (int i = 0; i < message.length(); i++) {
            char currentChar = message.charAt(i);
            if (currentChar == ' ') {
                continue;
            }

            int index1 = searchElement(Alphabet, currentChar);
            int index2 = searchNumber(SlowRight, SlowLeft[index1]);
            int index3 = searchNumber(MidRight, MidLeft[index2]);
            int index4 = searchNumber(FastRight, FastLeft[index3]);

            encryptedMessage.append(Alphabet[index4]);

            count1++;
            rotateWheel(FastLeft);
            rotateWheel(FastRight);

            if (count1 % 26 == 0 && count1 != 0) {
                rotateWheel(MidLeft);
                rotateWheel(MidRight);
                count2++;

                if (count2 % 26 == 0 && count2 != 0) {
                    rotateWheel(SlowLeft);
                    rotateWheel(SlowRight);
                }
            }
        }
        String encrypted = encryptedMessage.toString();
        System.out.println("加密后的结果为：" + encrypted);
        return encrypted;
    }

    // 解密函数
    public static void decodeFunction(String ciphertext) {
        System.out.println("解密后的结果为：");
        int count1 = 0;
        int count2 = 0;

        for (int i = 0; i < ciphertext.length(); i++) {
            char currentChar = ciphertext.charAt(i);
            if (currentChar == ' ') {
                continue;
            }

            int index1 = searchElement(Alphabet, currentChar);
            int index2 = searchNumber(FastLeft, FastRight[index1]);
            int index3 = searchNumber(MidLeft, MidRight[index2]);
            int index4 = searchNumber(SlowLeft, SlowRight[index3]);

            if (i == 0 || ciphertext.charAt(i - 1) == ' ') {
                System.out.print(Alphabet[index4]);
            } else {
                System.out.print(Character.toLowerCase(Alphabet[index4]));
            }

            count1++;
            rotateWheel(FastLeft);
            rotateWheel(FastRight);

            if (count1 % 26 == 0 && count1 != 0) {
                rotateWheel(MidLeft);
                rotateWheel(MidRight);
                count2++;

                if (count2 % 26 == 0 && count2 != 0) {
                    rotateWheel(SlowLeft);
                    rotateWheel(SlowRight);
                }
            }
        }
        System.out.println();
    }

    // 统计字母出现频率
    public static Map<Character, Double> countFrequency(String str) {
        HashMap<Character, Integer> map = new HashMap<>();
        for (int i = 0; i < str.length(); i++) {
            char ch = str.charAt(i);
            if (ch != ' ') {
                if (!map.containsKey(ch)) {
                    map.put(ch, 1);
                } else {
                    int val = map.get(ch);
                    map.replace(ch, val, val + 1);
                }
            }
        }

        Map<Character, Double> frequencyMap = new HashMap<>();
        for (Map.Entry<Character, Integer> entry : map.entrySet()) {
            frequencyMap.put(entry.getKey(), entry.getValue() * 1.0 / str.replace(" ", "").length());
        }
        return frequencyMap;
    }

    // 统计各字母加密前后频率
    public static void compareFrequency(String plaintext, String ciphertext) {
        Map<Character, Double> plaintextFrequency = countFrequency(plaintext);
        Map<Character, Double> ciphertextFrequency = countFrequency(ciphertext);

        for (char c : Alphabet) {
            double plainFreq = plaintextFrequency.getOrDefault(c, 0.0);
            double cipherFreq = ciphertextFrequency.getOrDefault(c, 0.0);
            System.out.printf("%c : 明文频率 %.4f, 密文频率 %.4f\n", c, plainFreq, cipherFreq);
        }
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        while (true) {
            System.out.println("请输入你要进行的操作：加密输入0，解密输入1");
            int operation = scanner.nextInt();
            scanner.nextLine();  // 清除换行符

            switch (operation) {
                case 0:
                    System.out.println("请输入想要加密的明文：");
                    String message = scanner.nextLine().toUpperCase();
                    String encryptedMessage = encryptFunction(message);
                    countFrequency(message);
                    countFrequency(encryptedMessage);
                    System.out.println("各字母加密前后频率对比：");
                    compareFrequency(message, encryptedMessage);
                    break;
                case 1:
                    System.out.println("请输入想要解密的密文：");
                    String ciphertext = scanner.nextLine().toUpperCase();
                    decodeFunction(ciphertext);
                    break;
                default:
                    System.out.println("请重新输入！");
                    continue;
            }
            break; // 结束循环
        }
        scanner.close();
    }
}